use std::{error::Error, process};
use csv::Reader;
use rusty_machine::{learning::{lin_reg::LinRegressor, SupModel}, linalg::{Matrix, Vector}};

fn machine_learning(path: &str) -> Result<LinRegressor, Box<dyn Error>> {
    let mut marked_vec = vec![];
    let mut target_vec = vec![];
    let mut reader = Reader::from_path(path).unwrap();
    let mut cnt = 0;
    for result in reader.records() {
        let record = result?;
        for sub_record in &record {
            if sub_record.eq("Iris-setosa") {
                marked_vec.push(0.);
            } else if sub_record.eq("Iris-versicolor") {
                marked_vec.push(1.);
            } else if sub_record.eq("Iris-virginica") {
                marked_vec.push(2.);
            } else {
                let tmp: f64 = sub_record.parse().unwrap();
                target_vec.push(tmp);
            }
        }
        cnt = cnt + 1;
    }
    let target_matrix = Matrix::new(cnt, 4, target_vec);
    let marked_matrix = Vector::new(marked_vec);
    
    let mut lr = LinRegressor::default();
    lr.train(&target_matrix, &marked_matrix).unwrap();
    
    Ok(lr)
}
fn main() {
    let dataset_path = "dataset/iris.data";

    match machine_learning(dataset_path) {
        Ok(lr) => {
            let test_matrix = Matrix::new(2, 4, vec![5.4, 2.9, 5.3, 2.2, 5.6, 2.7, 5.4, 2.3]);
            let output = lr.predict(&test_matrix);
            println!("{:#?}", output);
        },
        Err(err) => {
            println!("error running example: {}", err);
            process::exit(1);
        }
    }
}
